#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@author:hengk
@contact: hengk@foxmail.com
@datetime:2020-05-04 7:29
"""
from torchvision import transforms
from skimage import exposure
from skimage import util
from easydict import EasyDict
import numpy as np
import cv2
import random
import os
import torch
import yaml


class DataAugment:
    def __init__(self):
        pass
    def scale(self,img,longer_size):
        """
        对图片进行缩放
        :param longer_size:
        :return:
        """
        h,w,c = img.shape
        src_longer = max(h,w)
        s = longer_size/src_longer
        """
          如果放大图片，采用INTER_CUBIC方法
          如果缩小图片，采用INTER_AREA方法
        """
        if src_longer > 1:
            img = cv2.resize(img,(0,0),fx=s,fy=s,interpolation=cv2.INTER_CUBIC)
        else:
            img = cv2.resize(img,(0,0),fx=s,fy=s,interpolation=cv2.INTER_AREA)
        return img

    def crop(self,img,crop_w_ratio,crop_h_ratio):
        """
        对图片的长和宽进行裁边
        :param crop_w_ratio:
        :param crop_h_ratio:
        :return:
        """
        h,w,c = img.shape
        crop_h = int(h * crop_h_ratio)
        crop_w = int(w * crop_w_ratio)
        if random.random()>0.5:
            img = img[crop_h:h,crop_w:w,:]
        else:
            img = img[0:h-crop_h,0:w-crop_w,:]
        return img

    def change_light(self,img):
        """
        改变图片的光照强度
        :param image: 
        :return: 
        """
        flag = random.uniform(0.5, 1.5)
        #flag>1表示亮度变暗，flag<1表示亮度变亮
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = exposure.adjust_gamma(img,flag)
        img = cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
        return img
    def add_noise(self,img):
        """
        给图片增加噪音
        :param image:
        :return:
        """
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = util.random_noise(img, mode='gaussian') * 255
        img = img.astype(np.float32)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        return img

class DataLoader:
    def __init__(self,cfg):
        self.batch_size = cfg.train.batch_size
        self.label_path = os.path.join(cfg.train.data_dir,"labels.txt")
        self.images_dir = os.path.join(cfg.train.data_dir,"images")
        self.image_names = [line.split(" ")[0] for line in open(self.label_path)]
        self.labels = [line.split(" ")[1] for line in open(self.label_path)]
        self.len = len(self.labels)

        self.da = DataAugment()
        self.longer_size = cfg.train.longer_size
        self.crop_ratio_w = cfg.train.crop_ratio_w
        self.crop_ratio_h = cfg.train.crop_ratio_h
        self.trans = transforms.ToTensor()

    def get_len(self):
        return self.len

    def get_next(self):
        """
        获取下一组训练样本，每一次采样batchsize个样本，
        然后每一个样本进行扩增2倍，最后输出的样本个数为2*batchsize
        :return:
        """
        batch_idxs = random.sample(range(len(self.image_names)),self.batch_size)
        images = []
        labels = []
        for idx in batch_idxs:
            src_img  = cv2.imread(os.path.join(self.images_dir,self.image_names[idx]))

            for i in range(2):
                img = src_img
                #改变光照
                if random.random() > 0.7:
                    img = self.da.change_light(src_img)
                if random.random()>0.5:
                    #添加高斯噪音
                    img = self.da.add_noise(img)
                #对图片进行裁剪
                img = self.da.crop(img,self.crop_ratio_w,self.crop_ratio_h)
                #对图片进行归一化
                img = self.da.scale(img,self.longer_size)
                images.append(img)
                labels.append(self.labels[idx])
        #将数据加载到tensor中
        max_h = max([img.shape[0] for img in images])
        max_w = max([img.shape[1] for img in images])

        imgs_tensor = torch.zeros(size=(self.batch_size*2,3,max_h,max_w))
        labels_tensor = torch.zeros(size=(self.batch_size*2,1))

        for idx,(img,label) in enumerate(zip(images,labels)):
            img = self.trans(img)
            imgs_tensor[idx,:,0:img.shape[1],0:img.shape[2]] =img
            labels_tensor[idx] = int(label)
        # print(labels_tensor)
        return imgs_tensor,labels_tensor

if __name__ == '__main__':
    f = open("config/default.yaml", 'r', encoding='utf-8')
    cfg = yaml.load(f.read(),Loader=yaml.FullLoader)
    cfg = EasyDict(cfg)
    dl = DataLoader(cfg)
    imgs,labels = dl.get_next()
    print(imgs.shape,labels.shape)